package com.alibaba.halo.sandbox.app.context;

import com.alibaba.ttl.TransmittableThreadLocal;

import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * @author tieqin.lxy
 * @date 2020/05/08
 */
public class InvokeContext {

    private final Map<String, String> attachments = new ConcurrentHashMap<>();

    // @formatter:off
    /**
     * 使用 TransmittableThreadLocal
     */
    private static final TransmittableThreadLocal<InvokeContext> CONTEXT_THREAD_LOCAL = new TransmittableThreadLocal<InvokeContext>() {
        @Override
        protected InvokeContext initialValue() {
            return new InvokeContext();
        }
    };
    // @formatter:on

    /**
     * get context
     *
     * @return context
     */
    public static InvokeContext getContext() {
        return CONTEXT_THREAD_LOCAL.get();
    }

    /**
     * remove context
     */
    public static void removeContext() {
        CONTEXT_THREAD_LOCAL.remove();
    }

    /**
     * get attachment
     *
     * @param key
     * @return attachment
     */
    public String getAttachment(String key) {
        return attachments.get(key);
    }

    /**
     * set attachment
     *
     * @param key
     * @param value
     * @return context
     */
    public InvokeContext setAttachment(String key, String value) {
        if (value == null) {
            attachments.remove(key);
        } else {
            attachments.put(key, value);
        }
        return this;
    }

    /**
     * remove attachment
     *
     * @param key
     * @return context
     */
    public InvokeContext removeAttachment(String key) {
        attachments.remove(key);
        return this;
    }

    /**
     * get attachments
     *
     * @return attachments
     */
    public Map<String, String> getAttachments() {
        return attachments;
    }

    /**
     * set attachments
     *
     * @param attachment
     * @return context
     */
    public InvokeContext setAttachments(Map<String, String> attachment) {
        if (this.attachments == attachment) {
            return this;
        }
        this.attachments.clear();
        if (attachment != null && attachment.size() > 0) {
            this.attachments.putAll(attachment);
        }
        return this;
    }

    public void clearAttachments() {
        this.attachments.clear();
    }

}
